package optimizer;

import java.util.ArrayList;
import java.util.List;

import llvmir.value.Value;
import llvmir.value.constants.ConstantInt;
import llvmir.value.instructions.BinaryOp;
import llvmir.value.structure.BasicBlock;
import llvmir.value.structure.Function;
import llvmir.value.structure.Instruction;
import llvmir.value.structure.Module;
import mips.value.Immediate;
import mips.value.MipsValue;

public class ArithmeticOptimizer {
    // 乘除模优化
    // 默认关闭，在Optimizer中打开
    // 直接遍历module中的每一个指令
    public static boolean OPEN = false;

    public static void optimize() {
        // 遍历每一个指令
        // 如果是乘除模指令，且两个操作数都是常量
        // 则直接计算出结果，替换原来的指令
        Module module = Module.getInstance();
        for (Value f : module.getFunctions()) {
            // 为什么要优化三次呢，因为先模，模后可能还有乘
            optimizeFunction(f);
            optimizeFunction(f);
            optimizeFunction(f);
        }
        optimizeFunction(module.getMainFunction());
        optimizeFunction(module.getMainFunction());
        optimizeFunction(module.getMainFunction());
    }

    public static void optimizeFunction(Value f) {
        for (BasicBlock bb : new ArrayList<>(((Function) f).getBasicBlocks())) {
            for (Instruction inst : new ArrayList<>(bb.getInstructions())) {
                if (!(inst instanceof BinaryOp))
                    continue;
                BinaryOp binaryOp = (BinaryOp) inst;
                Value left = binaryOp.getOperands().get(0);
                Value right = binaryOp.getOperands().get(1);

                switch (binaryOp.getOp()) {
                    case "*":
                        optimizeMultiplication(bb, binaryOp, left, right);
                        break;
                    case "/":
                        optimizeDivision(bb, binaryOp, left, right);
                        break;
                    case "%":
                        optimizeModulo(bb, binaryOp, left, right);
                        break;
                    default:
                        break;
                }
            }
        }
    }

    private static void optimizeMultiplication(BasicBlock bb, BinaryOp binaryOp, Value left, Value right) {
        if (left instanceof ConstantInt) {
            optimizeConstantMultiplication(bb, binaryOp, (ConstantInt) left, right);
        } else if (right instanceof ConstantInt) {
            optimizeConstantMultiplication(bb, binaryOp, (ConstantInt) right, left);
        }
    }

    private static void optimizeConstantMultiplication(BasicBlock bb, BinaryOp binaryOp, ConstantInt constant,
            Value variable) {
        int value = constant.getValue();
        if (value == 0) {
            replaceInstructionWithZero(bb, binaryOp);
        } else if (value == 1) {
            replaceInstructionWithSelf(bb, binaryOp, variable);
        } else if (value > 0) {
            optimizePositivePowerOfTwo(bb, binaryOp, value, variable);
        } else {
            optimizeNegativePowerOfTwo(bb, binaryOp, value, variable);
        }
    }

    private static void optimizePositivePowerOfTwo(BasicBlock bb, BinaryOp binaryOp, int value, Value variable) {
        int powerOfTwo = Integer.highestOneBit(value);
        if (powerOfTwo == value) {
            // 2^n
            createShiftInstruction(bb, binaryOp, variable, powerOfTwo);
        } else if (powerOfTwo == value + 1 || powerOfTwo == value - 1) {
            // 2^n - 1 or 2^n + 1
            createShiftOpInstruction(bb, binaryOp, variable, powerOfTwo, value > powerOfTwo);
        }
    }

    private static void optimizeNegativePowerOfTwo(BasicBlock bb, BinaryOp binaryOp, int value, Value variable) {
        int absValue = Math.abs(value);
        int powerOfTwo = Integer.highestOneBit(absValue);
        if (powerOfTwo == absValue) {
            createNegativeShiftInstruction(bb, binaryOp, variable, powerOfTwo);
        }
    }

    private static void optimizeDivision(BasicBlock bb, BinaryOp binaryOp, Value left, Value right) {
        if (right instanceof ConstantInt) {
            int value = ((ConstantInt) right).getValue();
            if (value == 1) {
                replaceInstructionWithSelf(bb, binaryOp, left);
            } else if (value == -1) {
                replaceInstructionWithNegation(bb, binaryOp, left);
            }
        }
    }

    private static void optimizeModulo(BasicBlock bb, BinaryOp binaryOp, Value left, Value right) {
        if (right instanceof ConstantInt) {
            int value = ((ConstantInt) right).getValue();
            if (Math.abs(value) == 1) {
                replaceInstructionWithZero(bb, binaryOp);
            } else {
                createModuloInstruction(bb, binaryOp, left, right);
            }
        }
    }

    private static void replaceInstructionWithZero(BasicBlock bb, BinaryOp binaryOp) {
        BinaryOp newInst = new BinaryOp(binaryOp.getOperands().get(0), binaryOp.getOperands().get(0), "-",
                binaryOp.getBasicBlock());
        bb.replaceInstruction(binaryOp, List.of(newInst));
    }

    private static void replaceInstructionWithSelf(BasicBlock bb, BinaryOp binaryOp, Value self) {
        BinaryOp newInst = new BinaryOp(self, new ConstantInt(0), "+", binaryOp.getBasicBlock());
        bb.replaceInstruction(binaryOp, List.of(newInst));
    }

    private static void replaceInstructionWithNegation(BasicBlock bb, BinaryOp binaryOp, Value variable) {
        BinaryOp newInst = new BinaryOp(new ConstantInt(0), variable, "-", binaryOp.getBasicBlock());
        bb.replaceInstruction(binaryOp, List.of(newInst));
    }

    private static void createShiftInstruction(BasicBlock bb, BinaryOp binaryOp, Value variable, int powerOfTwo) {
        int exponent = Integer.numberOfTrailingZeros(powerOfTwo);
        BinaryOp newInst = new BinaryOp(variable, new ConstantInt(exponent), "<<", binaryOp.getBasicBlock());
        bb.replaceInstruction(binaryOp, List.of(newInst));
    }

    private static void createShiftOpInstruction(BasicBlock bb, BinaryOp binaryOp, Value variable,
            int powerOfTwo, boolean isAdd) {
        int exponent = Integer.numberOfTrailingZeros(powerOfTwo);
        BinaryOp shiftInst = new BinaryOp(variable, new ConstantInt(exponent), "<<", binaryOp.getBasicBlock());
        BinaryOp addOrSubInst = new BinaryOp(shiftInst, variable, isAdd ? "+" : "-", binaryOp.getBasicBlock());
        bb.replaceInstruction(binaryOp, List.of(shiftInst, addOrSubInst));
    }

    private static void createNegativeShiftInstruction(BasicBlock bb, BinaryOp binaryOp, Value variable,
            int powerOfTwo) {
        int exponent = Integer.numberOfTrailingZeros(powerOfTwo);
        BinaryOp shiftInst = new BinaryOp(variable, new ConstantInt(exponent), "<<", binaryOp.getBasicBlock());
        BinaryOp negInst = new BinaryOp(new ConstantInt(0), shiftInst, "-", binaryOp.getBasicBlock());
        bb.replaceInstruction(binaryOp, List.of(shiftInst, negInst));
    }

    private static void createModuloInstruction(BasicBlock bb, BinaryOp binaryOp, Value left, Value right) {
        BinaryOp divInst = new BinaryOp(left, right, "/", binaryOp.getBasicBlock());
        BinaryOp mulInst = new BinaryOp(divInst, right, "*", binaryOp.getBasicBlock());
        BinaryOp subInst = new BinaryOp(left, mulInst, "-", binaryOp.getBasicBlock());
        bb.replaceInstruction(binaryOp, List.of(divInst, mulInst, subInst));
    }

    // 后端除法优化
    public static String optimizeDivision(MipsValue op1, MipsValue op2, MipsValue dst, String type) {
        StringBuilder sb = new StringBuilder();

        int N = 31;
        long d = ((Immediate) op2).getValue();
        long absD = Math.abs(d);
        long l = Math.max(1, log2x(absD));
        long m = 0;
        // 依次枚举 ℓ，验证 2 ^{N+ℓ} ≤ m × d ≤ 2 ^{N+ℓ} + 2^ℓ 是否成立，得到 一组合法的 m, ℓ。
        boolean isVaild = false;
        for (; l <= 31; l++) {
            long left = 1L << (N + l);
            long right = left + (1L << l);
            m = (left + absD - 1) / absD;
            if (m * absD >= left && m * absD <= right) {
                isVaild = true;
                break;
            }
        }
        if (isVaild) {
            /*
             * $n/d=
             * \begin{cases}n ,if(d=1)\\
             * SRA(n+SRL+(SRA(n,k-1),32-k),k),if(d=2^k)\\
             * SRA(MULSH(m,n),ℓ-1)+SRL(n ,31),if(m<2^{31})\\
             * SRA(n+MULSH(m-2^{32},n),ℓ-1)+SRL(n,31),otherwise
             * \end{cases} $
             */
            if (absD == 1) {
                sb.append("move ").append(dst).append(", ").append(op1).append("\n");
            } else if ((absD & absD - 1) == 0) {
                // d = 2^k
                int k = log2x(absD);
                sb.append("sra $v0, ").append(op1).append(", ").append(k - 1).append("\n");
                sb.append("srl $v0, $v0, ").append(32 - k).append("\n");
                sb.append("add $v0, $v0, ").append(op1).append("\n");
                sb.append("sra ").append(dst).append(", $v0, ").append(k).append("\n");
                if (d < 0) {
                    sb.append("sub ").append(dst).append(", $zero, ").append(dst).append("\n");
                }
            } else if (m < (1L << 31)) {
                sb.append("li $v0, ").append(m).append("\n");
                sb.append("mult ").append(op1).append(", $v0\n");
                sb.append("mfhi $v0\n");
                sb.append("sra $v0, $v0, ").append(l - 1).append("\n");
                sb.append("srl ").append(dst).append(", ").append(op1).append(", 31\n");
                sb.append("addu ").append(dst).append(", $v0, ").append(dst).append("\n");
                if (d < 0) {
                    sb.append("sub ").append(dst).append(", $zero, ").append(dst).append("\n");
                }
            } else {
                long m2 = m - (1L << 32);
                sb.append("li $v0, ").append(m2).append("\n");
                sb.append("mult ").append(op1).append(", $v0\n");
                sb.append("mfhi $v0\n");
                sb.append("addu $v0, $v0, ").append(op1).append("\n");
                sb.append("sra $v0, $v0, ").append(l - 1).append("\n");

                sb.append("srl ").append(dst).append(", ").append(op1).append(", 31\n");
                sb.append("addu ").append(dst).append(", $v0, ").append(dst).append("\n");
                if (d < 0) {
                    sb.append("sub ").append(dst).append(", $zero, ").append(dst).append("\n");
                }
            }
            return sb.toString();
        }

        if (type.equals("div")) {
            return "li $v0, " + op2 + "\ndiv " + op1 + ", $v0\nmflo " + dst;
        } else {
            return "li $v0, " + op2 + "\ndiv " + op1 + ", $v0\nmfhi " + dst;
        }

    }

    public static int log2x(long x) {
        int res = 0;
        while (x > 1) {
            x >>= 1;
            res++;
        }
        return res;
    }
}
